package middleware

import (
	"bytes"
	"context"
	"io"
	"time"

	"gitee.com/chewel/hermes/pkg/journal"
	"gitee.com/chewel/hermes/pkg/utils"
	"github.com/gin-gonic/gin"
	"go.uber.org/atomic"
)

type AccessRecord struct {
	TraceId    string
	Method     string
	URL        string
	ReqBody    string
	RespBody   string
	Duration   time.Duration
	StatusCode int
	Error      string
	ClientIP   string
	UserAgent  string
}

type LoggerBuilder struct {
	allowReqBody  *atomic.Bool
	allowRespBody *atomic.Bool
	logger        journal.ILogger
	loggerFunc    func(ctx context.Context, logger journal.ILogger, record *AccessRecord)
}

func NewLoggerBuilder(logger journal.ILogger) *LoggerBuilder {
	return &LoggerBuilder{
		logger:        logger,
		allowReqBody:  atomic.NewBool(false),
		allowRespBody: atomic.NewBool(false),
	}
}

func (b *LoggerBuilder) AllowReqBody(allowed bool) *LoggerBuilder {
	b.allowReqBody.Store(allowed)
	return b
}

func (b *LoggerBuilder) AllowRespBody(allowed bool) *LoggerBuilder {
	b.allowRespBody.Store(allowed)
	return b
}

func (b *LoggerBuilder) LoggerFunc(fn func(ctx context.Context, logger journal.ILogger, record *AccessRecord)) *LoggerBuilder {
	b.loggerFunc = fn
	return b
}

func (b *LoggerBuilder) Build() gin.HandlerFunc {
	return func(c *gin.Context) {
		start := time.Now()
		traceId := c.GetString("X-Request-ID")
		if traceId == "" {
			traceId = utils.NewID()
		}
		record := &AccessRecord{
			TraceId:   traceId,
			Method:    c.Request.Method,
			URL:       c.Request.URL.String(),
			ClientIP:  c.ClientIP(),
			UserAgent: c.Request.UserAgent(),
		}
		if b.allowReqBody.Load() && c.Request.Body != nil {
			body, err := c.GetRawData()
			if err == nil {
				record.ReqBody = string(body)
				c.Request.Body = io.NopCloser(bytes.NewReader(body))
				c.Request.GetBody = func() (io.ReadCloser, error) {
					return c.Request.Body, nil
				}
			}
		}
		if b.allowRespBody.Load() {
			c.Writer = &responseWriter{
				ResponseWriter: c.Writer,
				record:         record,
			}
		}

		defer func() {
			record.StatusCode = c.Writer.Status()
			record.Duration = time.Since(start)
			record.Error = c.Errors.ByType(gin.ErrorTypePrivate).String()

			if b.loggerFunc != nil {
				b.loggerFunc(c, b.logger, record)
			} else {
				b.logger.Debug("[QGM]",
					journal.String("trace_id", record.TraceId),
					journal.Duration("duration", record.Duration),
					journal.Int("status", record.StatusCode),
					journal.String("method", record.Method),
					journal.String("url", record.URL),
					journal.String("client_ip", record.ClientIP),
					journal.String("user_agent", record.UserAgent),
					journal.String("req_body", record.ReqBody),
					journal.String("res_body", record.RespBody),
					journal.String("errors", record.Error),
				)
			}
		}()
		c.Next()
	}
}

type responseWriter struct {
	gin.ResponseWriter
	record *AccessRecord
}

func (r *responseWriter) Write(data []byte) (int, error) {
	r.record.RespBody = string(data)
	return r.ResponseWriter.Write(data)
}

func (r *responseWriter) WriteHeader(statusCode int) {
	r.record.StatusCode = statusCode
	r.ResponseWriter.WriteHeader(statusCode)
}

func (r *responseWriter) WriteString(data string) (int, error) {
	r.record.RespBody = data
	return r.ResponseWriter.WriteString(data)
}
